// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <tbir/runtime/container.h>
#include <tbir/runtime/threadpool/i_thread_pool.h>

#include <memory>
#include <unordered_set>

namespace tbir::runtime {

    class ThreadPoolExecutor {
    public:
        explicit ThreadPoolExecutor(const std::shared_ptr<internal::IThreadPool> &pool, bool lock_free);

        ~ThreadPoolExecutor() = default;

        List ParallelFor(const UserDataRef &op,
                         const List &inputs,
                         int64_t expt_num_threads,
                         int64_t group_size);

        List ParallelFor(const UserDataRef &op, const List &inputs);

        Tuple ParallelFor(const UserDataRef &op,
                          const Tuple &inputs,
                          int64_t expt_num_threads,
                          int64_t group_size);

        Tuple ParallelFor(const UserDataRef &op, const Tuple &inputs);

        List ParallelStarMap(const UserDataRef &op,
                             const List &inputs,
                             int64_t expt_num_threads,
                             int64_t group_size);

        List ParallelStarMap(const UserDataRef &op, const List &inputs);

        Tuple ParallelStarMap(const UserDataRef &op,
                              const Tuple &inputs,
                              int64_t expt_num_threads,
                              int64_t group_size);

        Tuple ParallelStarMap(const UserDataRef &op, const Tuple &inputs);

        RTValue ApplyAsync(const UserDataRef &op, const PyArgs &args);

        RTValue Submit(PyArgs args);

    private:
        void ParallelForImpl(const UserDataRef &op,
                             const Any *inputs_begin,
                             const Any *inputs_end,
                             int64_t expt_num_threads,
                             int64_t group_size,
                             RTValue *outputs_begin,
                             bool unpack_args);

    private:
        bool lock_free_ = true;
        int thread_num_ = 0;
        std::shared_ptr<internal::IThreadPool> pool_ = nullptr;
        std::atomic<size_t> serial_{0};
        std::unordered_set<std::thread::id> pool_thread_ids_;
    };

}  // namespace tbir::runtime
